package org.ainlolcat.ainscience.statistics.handlers.fitting;

import org.ainlolcat.ainscience.statistics.handlers.fitting.functions.ZeroBasedGaussFunction;
import org.ainlolcat.ainscience.statistics.handlers.fitting.functions.ZeroBasedMultipleGaussFunction;
import org.ainlolcat.ainscience.statistics.handlers.fitting.functions.ZeroBasedPsdVoigtFunction;
import org.ainlolcat.ainscience.utils.ArraysUtil;
import org.apache.commons.math.optimization.fitting.CurveFitter;
import org.apache.commons.math.optimization.fitting.ParametricRealFunction;
import org.apache.commons.math.optimization.general.LevenbergMarquardtOptimizer;
import org.apache.log4j.Logger;

import java.util.ArrayList;
import java.util.List;

/**
* @author ainlolcat
*         Date: 5/9/13
*/
public class MultipleGaussPeakFitting {
    private static Logger log = Logger.getLogger(MultipleGaussPeakFitting.class);

    ParametricRealFunction function;
    double mean;
    double disp;
    double maximum;
    private MultiplePeakFittingStatisticHandler peakFittingStatisticHandler;

    public MultipleGaussPeakFitting(MultiplePeakFittingStatisticHandler peakFittingStatisticHandler) {
        this.peakFittingStatisticHandler = peakFittingStatisticHandler;

    }

    public List<Peak> getParamsFor(double[] x, double[] y, List<Peak> peaks,  String name) {
        //prepare graph for fitting. normalize. search for peak center and other
        function = new ZeroBasedMultipleGaussFunction(peaks);
        double[] y1 = prepareGraph(x,y);

        double[][] temp = new double[3][y1.length];
        temp[1] = y1;

        LevenbergMarquardtOptimizer optimizer = new LevenbergMarquardtOptimizer();
        //fill data points
        CurveFitter fitter = new CurveFitter(optimizer);
        for (int i=0;i<(x.length<y.length?x.length:y.length);i++){
            fitter.addObservedPoint(x[i], y1[i]);
        }


        //lets try
        try {
            //prepare initial guess
            List<Double> initial = new ArrayList<Double>();
            for (int i=0;i<peaks.size();i++){
                initial.add(peaks.get(i).A);
                if (!peaks.get(i).fixed)
                    initial.add(peaks.get(i).x0);
                initial.add(peaks.get(i).w);
            }

            double[] initialguess1 = ArraysUtil.toArray(initial);
            int[] shifts;
            shifts = new int[peaks.size()];
            shifts[0] = peaks.get(0).fixed ? 1 : 0;
            for (int i=1;i<peaks.size();i++){
                shifts[i] = shifts[i-1] +  (peaks.get(i).fixed ? 1 : 0);
            }

            double divA = function.value(mean,initialguess1);
            for (int i = 0; i< peaks.size(); i++){
                int shift = i > 0 ? shifts[i-1] :  0;
                initialguess1[i * 3 + 0 - shift] = initialguess1[i * 3 + 0 - shift] / divA;
            }

            double[] result = fitter.fit(function, initialguess1);

            for (int i=0;i<y1.length;i++){
                temp[0][i] = x[i];
                temp[2][i]= function.value(x[i], result);
            }

            peakFittingStatisticHandler.data.put( name , temp);

            List<Peak> answer = new ArrayList<Peak>();
            for (int i=0;i<peaks.size();i++){
                Peak peak = new Peak();
                Peak current = peaks.get(i);
                int shift = i > 0 ? shifts[i-1] :  0;
                peak.A = result[i * 3 + 0 - shift];
                peak.x0 = current.fixed ? current.x0 : result[i * 3 + 1 - shift];
                peak.w =  current.fixed ? result[i * 3 + 1 - shift] : result[i * 3 + 2 - shift];
                answer.add(peak);
            }

            return answer;
        } catch (Exception e) {
            log.error(e);
            return null;
        }
    }

    private double[] prepareGraph(double[] x, double[] y){
        int baseLineDots = peakFittingStatisticHandler.baseline;
        if (baseLineDots > x.length/6){
            baseLineDots = x.length/6;
        }
        maximum = y[0];
        mean = x[0];
        double ySumm = 0;
        double y0 = 0,x0 = 0,y1 = 0,x1 = 0;
        for (int i=0;i<baseLineDots;i++){
            x0+=x[i];
            y0+=y[i];
            x1+=x[x.length-1-i];
            y1+=y[y.length-1-i];
        }
        x0 = x0/baseLineDots;
        y0 = y0/baseLineDots;
        x1 = x1/baseLineDots;
        y1 = y1/baseLineDots;
        double k = (y1-y0)/(x1-x0);
        double[] out = new double[x.length<y.length?x.length:y.length];
        for (int i=0;i<(x.length<y.length?x.length:y.length);i++){
            out[i] = (y[i] - (y0 + k * (x[i] - x0)));
        }
        for (int i=0;i<(x.length<y.length?x.length:y.length);i++){
            maximum = y[i]>maximum?y[i]:maximum;
            ySumm+=y[i];
            mean += y[i]*x[i];
            disp += y[i]*x[i]*x[i];
        }
        mean = mean / ySumm;
        disp = disp / ySumm - mean * mean;
        for (int i=0;i<(x.length<y.length?x.length:y.length);i++){
            out[i] = out[i]/maximum;
        }
        return out;
    }
}
